#pragma once

#include <functional>
#include <ostream>
#include <string>
#include <unordered_map>
#include <vector>

#include "arithmetic.h"
#include "arm.h"
#include "ast.h"
#include "memobj.h"
#include "utils.h"

struct ArmBlock;
struct ArmFunction;

namespace Arm {

    struct Inst : public Printable {
        ArmBlock* parent;
    };

    // add, sub, sdiv, mul (I)
    // vadd.f32, vsub.f32, vdiv.f32, vmul.f32 (F)
    struct BinaryOp : public Inst {
        Register rd;
        Register rs, rt;
        ArithOp op;

        BinaryOp(Register rd, Register rs, Register rt, ArithOp op);

        void GenAsm(std::ostream& os) const override;
    };

    // cmp (I), vcmp (F)
    struct Cmp : public Inst {
        Register rs, rt;

        Cmp(Register rs, Register rt);

        void GenAsm(std::ostream& os) const override;
    };

    // add, sub (I)
    struct BinaryOpImm : public Inst {
        Register rd;
        Register rs;
        int rt;
        ArithOp op;

        BinaryOpImm(Register rd, Register rs, int rt, ArithOp op);

        void GenAsm(std::ostream& os) const override;
    };

    // vcvt.s32.f32 (F->I), vcvt.f32.s32 (I->F)
    // mov (+I), vmov.f32 (+F)
    // rsb rd, rs, #0 (-I), vneg.f32 (-F)
    struct UnaryOp : public Inst {
        Register rd;
        Register rs;
        ArithOp op;

        UnaryOp(Register rd, Register rs, ArithOp op);

        void GenAsm(std::ostream& os) const override;
    };

    // mov (I->I), vmov.f32 (F->F)
    // vmov (I->F / F->I)
    struct Mov : public Inst {
        Register rd;
        Register rs;

        Mov(Register rd, Register rs);

        void GenAsm(std::ostream& os) const override;
    };

    // movw + movt (rd <- &var)
    struct LoadAddr : public Inst {
        Register rd;
        std::string name;

        LoadAddr(Register rd, const std::string& name);

        void GenAsm(std::ostream& os) const override;
    };

    // movw + movt (rd <- imm)
    struct LoadImm : public Inst {
        Register rd;
        int imm;

        LoadImm(Register rd, int imm);

        void GenAsm(std::ostream& os) const override;
    };

    // vldr.32 rd, label
    struct LoadFloat : public Inst {
        Register rd;
        std::string lbl;

        LoadFloat(Register rd, const std::string& lbl);

        void GenAsm(std::ostream& os) const override;
    };

    // rd <- [mem]
    struct Load : public Inst {
        Register rd;
        MemObj* mem;

        Load(Register rd, MemObj* mem);

        void GenAsm(std::ostream& os) const override;
    };

    // rs -> [mem]
    struct Store : public Inst {
        Register rs;
        MemObj* mem;

        Store(Register rs, MemObj* mem);

        void GenAsm(std::ostream& os) const override;
    };

    // bl (pc -> lr)
    struct Call : public Inst {
        std::string name;

        Call(const std::string& name);

        void GenAsm(std::ostream& os) const override;
    };

    struct Terminator : public Inst {
        virtual std::vector<ArmBlock*> GetNext() const = 0;
        virtual ArmBlock* GetDefaultJump() const = 0;

        void GenJumpAsm(std::ostream& os) const;
    };

    // b
    struct Jump : public Terminator {
        ArmBlock* nxt;

        Jump(ArmBlock* nxt);

        void GenAsm(std::ostream& os) const override;
        std::vector<ArmBlock*> GetNext() const override;
        ArmBlock* GetDefaultJump() const override;
    };

    // b{cond}
    struct JumpCond : public Terminator {
        ArithOp cond;
        ArmBlock* t;
        ArmBlock* f;

        JumpCond(ArithOp cond, ArmBlock* t, ArmBlock* f);

        void GenAsm(std::ostream& os) const override;
        std::vector<ArmBlock*> GetNext() const override;
        ArmBlock* GetDefaultJump() const override;
    };

    // pop lr / mov pc, lr
    struct Ret : public Terminator {
        void GenAsm(std::ostream& os) const override;
        std::vector<ArmBlock*> GetNext() const override;
        ArmBlock* GetDefaultJump() const override;
    };

    // Pseudo Instruction
    // push {...} (I), vpush {...} (F)
    struct Push : public Inst {
        Type::Basic type;
        std::vector<Register> regs;

        Push(Register reg);
        Push(const std::vector<Register>& regs);

        void GenAsm(std::ostream& os) const override;
    };

    // pop {...} (I), vpop {...} (F)
    struct Pop : public Inst {
        Type::Basic type;
        std::vector<Register> regs;

        Pop(Register reg);
        Pop(const std::vector<Register>& regs);

        void GenAsm(std::ostream& os) const override;
    };

}  // namespace Arm
